import cv2
import numpy as np
import matplotlib.pyplot as plt


def sobel_filter(image):
    """
    对输入图像应用Sobel算子滤波，分别计算水平和垂直方向梯度，然后求梯度幅值。

    Parameters:
    image (numpy.ndarray): 输入的灰度图像。

    Returns:
    numpy.ndarray: 经过Sobel算子滤波后的图像。
    """
    # Sobel算子的水平方向卷积核
    sobel_x = np.array([[-1, 0, 1],
                        [-2, 0, 2],
                        [-1, 0, 1]])
    # Sobel算子的垂直方向卷积核
    sobel_y = np.array([[-1, -2, -1],
                        [0, 0, 0],
                        [1, 2, 1]])
    # 获取图像的高度和宽度，对于彩色图像有3个通道所以取前两个维度
    height, width = image.shape[:2]
    # 创建一个与输入图像尺寸相同的全零数组用于存储结果
    new_image = np.zeros((height, width))
    # 对图像进行填充，填充宽度为 1 像素，填充值为常量 0，为卷积操作做准备
    padded_image = np.pad(image, ((1, 1), (1, 1)), 'constant')
    # 遍历图像的每个像素
    for i in range(height):
        for j in range(width):
            # 取 3x3 的邻域用于水平方向的卷积操作
            neighborhood_x = padded_image[i:i + 3, j:j + 3]
            # 计算水平方向的梯度
            gx = np.sum(neighborhood_x * sobel_x)
            # 取 3x3 的邻域用于垂直方向的卷积操作
            neighborhood_y = padded_image[i:i + 3, j:j + 3]
            # 计算垂直方向的梯度
            gy = np.sum(neighborhood_y * sobel_y)
            # 计算梯度幅值
            new_image[i, j] = np.sqrt(gx ** 2 + gy ** 2)
    return new_image


def custom_filter(image):
    """
    使用给定的卷积核对图像进行滤波操作。
    """
    # 自定义的卷积核
    kernel = np.array([[1, 0, -1],
                       [2, 0, -2],
                       [1, 0, -1]])
    # 获取卷积核的尺寸
    kernel_size = kernel.shape[0]
    # 获取图像的高度和宽度
    height, width = image.shape[:2]
    # 创建一个与输入图像尺寸相同的全零数组用于存储结果
    new_image = np.zeros((height, width))
    # 计算需要填充的尺寸
    pad_size = kernel_size // 2
    # 对图像进行填充
    padded_image = np.pad(image, ((pad_size, pad_size), (pad_size, pad_size)), 'constant')
    # 遍历图像的每个像素
    for i in range(height):
        for j in range(width):
            # 取与卷积核尺寸相同的邻域
            neighborhood = padded_image[i:i + kernel_size, j:j + kernel_size]
            # 进行卷积操作
            new_image[i, j] = np.sum(neighborhood * kernel)
    return new_image


def plot_color_histogram(image):
    """
    绘制输入彩色图像的颜色直方图（分别绘制R、G、B三个通道的直方图）。
    """
    # 颜色列表，分别对应红、绿、蓝三个通道
    colors = ('r', 'g', 'b')
    # 遍历颜色列表
    for i, color in enumerate(colors):
        # 绘制相应通道的直方图
        channel_hist = plt.hist(image[:, :, i].ravel(), bins=256, range=(0, 255),
                                fc=color, ec=color, alpha=0.5)
    # 设置 x 轴标签
    plt.xlabel('Pixel Intensity')
    # 设置 y 轴标签
    plt.ylabel('Frequency')
    # 设置图表标题
    plt.title('Color Histogram')
    # 显示直方图
    plt.show()


def extract_texture_features(image):
    """
    提取图像的纹理特征（这里以灰度共生矩阵对比度为例），返回纹理特征值。
    """
    # 定义距离 d
    d = 1
    # 定义角度 theta（这里为 0，表示水平方向）
    theta = 0
    # 获取图像的高度和宽度
    height, width = image.shape[:2]
    # 创建一个 256x256 的全零矩阵用于存储灰度共生矩阵
    glcm = np.zeros((256, 256))
    # 遍历图像像素，计算灰度共生矩阵
    for i in range(height - d):
        for j in range(width - d):
            row = image[i, j]
            col = image[i + d, j]
            glcm[row, col] += 1
    # 初始化对比度为 0
    contrast = 0
    # 计算灰度共生矩阵的对比度
    for i in range(256):
        for j in range(256):
            contrast += (i - j) ** 2 * glcm[i, j]
    return contrast


if __name__ == "__main__":
    # 使用 cv2 读取图像，参数 1 表示读取彩色图像
    image = cv2.imread('test1.jpeg', 1)
    # 将彩色图像转换为灰度图像，方便后续处理（若本身就是灰度图可省略这步），这里先保留彩色图像用于绘制颜色直方图
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 将图像转换为 RGB 格式，用于 matplotlib 显示
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 对灰度图像应用 Sobel 算子滤波
    sobel_result = sobel_filter(gray_image)
    # 对灰度图像使用自定义卷积核滤波
    custom_result = custom_filter(gray_image)

    # 展示原始图像
    plt.subplot(131)
    # 显示原始图像
    plt.imshow(image_rgb)
    # 设置子图标题
    plt.title('Original Image')
    # 关闭坐标轴显示
    plt.axis('off')

    # 展示 Sobel 算子滤波结果
    plt.subplot(132)
    # 显示 Sobel 滤波结果，使用灰度图显示
    plt.imshow(sobel_result, cmap='gray')
    # 设置子图标题
    plt.title('Sobel Filter Result')
    # 关闭坐标轴显示
    plt.axis('off')

    # 展示给定卷积核滤波结果
    plt.subplot(133)
    # 显示自定义卷积核滤波结果，使用灰度图显示
    plt.imshow(custom_result, cmap='gray')
    # 设置子图标题
    plt.title('Custom Filter Result')
    # 关闭坐标轴显示
    plt.axis('off')
    # 显示图像窗口
    plt.show()

    # 绘制原始彩色图像的直方图
    plot_color_histogram(image)

    # 提取纹理特征
    texture_feature = extract_texture_features(gray_image)
    # 将提取的纹理特征保存为 npy 文件
    np.save('E://MV//exp1//texture_feature.npy', texture_feature)

